import torch
import torch.nn as nn

class TransE(nn.Module):
    def __init__(self, num_entities, num_relations, hidden_dim):
        self._num_entities = num_entities
        self._num_relations = num_relations
        self._hidden_dim = hidden_dim

        super(TransE, self).__init__()

        self.entities_embedding = nn.Embedding(self._num_entities, self._hidden_dim)
        self.relations_embedding = nn.Embedding(self._num_relations, self._hidden_dim)

    def forward(self, inputs):
        head = self.entities_embedding(inputs[0])
        relation = self.relations_embedding(inputs[1])
        tail = self.entities_embedding(inputs[2])
        return self.norm_l2(head, relation, tail)
    
    def predict(self, inputs):
        """_summary_

        Args:
            inputs (_type_): _description_(head, relation). head: (B,C,D)

        Returns:
            _type_: _description_
        """
        # 不计算梯度
        with torch.no_grad():
            head = self.entities_embedding(inputs[0])[:,None,:].repeat(1, self._num_entities, 1)
            relation = self.relations_embedding(inputs[1])[:,None,:].repeat(1, self._num_entities, 1)
            tail_score = self.norm_l2(head, relation, self.entities_embedding.weight.data)
        return tail_score
    
    @staticmethod
    def norm_l2(h,r,t):
        d = h+r-t
        return torch.norm(torch.square(d), p=2, dim=d.dim()-1)  # 最后1个维度上计算L2范数
